/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph_optimizer/ub_fusion/tbe_pass/tbe_depthwiseconv_dequant_fusion_pass.h"
#include <string>
#include <vector>

namespace fe {
using std::vector;
static const char PATTERN_CONV[] = "DepthwiseConvolution";
static const char PATTERN_DEQUANT[] = "dequant";
static const char PATTERN_RELU[] = "relu";
static const char PATTERN_RELU6[] = "relu6";
static const char PATTERN_QUANT[] = "quant";
static const char PATTERN_REQUANT[] = "requant";
static const char PATTERN_OTHER_INPUT[] = "otherInput";
static const char PATTERN_OTHER_INPUT2[] = "otherInput2";
static const string PATTERN_MUL = "mul";
static const string PATTERN_SIGMOID = "sigmoid";
static const string PATTERN_POWER1 = "power1";
static const string PATTERN_POWER2 = "power2";
static const string PATTERN_ELEMWISE = "eltwise";
static const bool IGNORE_INPUT_NUM = true;
static const bool IGNORE_OUTPUT_NUM = true;
static const int TBE_OUTPUT_BRANCH_SINGE = 1;
/*
 * @brief:  define convolution and single input op fusion pattern
 *
 * pattern configuration limit:
 * 1. total min value must be 1 for all head candidated desc.
 * 2. any head candidated desc max value must be 1.
 * 3. output desc can not be itself.
 *
 * DepthwiseConvolution-->AcendDeQuant
 *
 * fusion node: AcendDeQuant,  DepthwiseConvolution
 *
 * @return BufferFusionPattern: return all valid patterns.
 */
vector<BufferFusionPattern *> TbeDepthwiseConvDequantFusionPass::DefinePatterns() {
  vector<BufferFusionPattern *> patterns;

  string pass_name = "TbeDepthwiseConvDequantReluFusion";
  BufferFusionPattern *pattern = new (std::nothrow) BufferFusionPattern(pass_name);
  FE_CHECK((pattern == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_RELU, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_QUANT, {OP_PATTERN_QUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_RELU})
      .SetOutputs(PATTERN_RELU, {PATTERN_QUANT})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern);
  FE_LOGD("End to define %s pass pattern.", pass_name.c_str());


  string pass_name2 = "TbeDepthwiseConvDequantMulQuantFusion";
  BufferFusionPattern *pattern2 = new (std::nothrow) BufferFusionPattern(pass_name2);
  FE_CHECK((pattern2 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name2.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern2->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_MUL, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT2, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_QUANT, {OP_PATTERN_QUANT}, TBE_PATTERN_NUM_NONE, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_MUL})
      .SetOutputs(PATTERN_OTHER_INPUT2, {PATTERN_MUL})
      .SetOutputs(PATTERN_MUL, {PATTERN_QUANT})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern2);
  FE_LOGD("End to define %s pass pattern.", pass_name2.c_str());

  string pass_name3 = "TbeDepthwiseConvDequantMulFusion";
  BufferFusionPattern *pattern3 = new (std::nothrow) BufferFusionPattern(pass_name3);
  FE_CHECK((pattern3 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name3.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern3->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_MUL, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT2, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_MUL})
      .SetOutputs(PATTERN_OTHER_INPUT2, {PATTERN_MUL})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern3);
  FE_LOGD("End to define %s pass pattern.", pass_name3.c_str());

  string pass_name4 = "TbeDepthwiseConvDequantSigmoidMulFusion";
  BufferFusionPattern *pattern4 = new (std::nothrow) BufferFusionPattern(pass_name4);
  FE_CHECK((pattern4 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name4.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern4->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_SIGMOID, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_MUL, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_SIGMOID, PATTERN_MUL}, TBE_OUTPUT_BRANCH_MULTI)
      .SetOutputs(PATTERN_SIGMOID, {PATTERN_MUL})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern4);
  FE_LOGD("End to define %s pass pattern.", pass_name4.c_str());

  string pass_name1 = "TbeDepthwiseConvDequantFusion";
  BufferFusionPattern *pattern1 = new (std::nothrow) BufferFusionPattern(pass_name1);
  FE_CHECK((pattern1 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name1.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern1->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_QUANT, {OP_PATTERN_QUANT}, TBE_PATTERN_NUM_NONE, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_QUANT})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern1);
  FE_LOGD("End to define %s pass pattern.", pass_name1.c_str());

  string pass_name5 = "TbeDepthwiseConvRequantFusion";
  BufferFusionPattern *pattern5 = new (std::nothrow) BufferFusionPattern(pass_name5);
  FE_CHECK((pattern5 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name5.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern5->AddOpDesc(PATTERN_REQUANT, {OP_PATTERN_REQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_REQUANT})
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_REQUANT});
  patterns.push_back(pattern5);
  FE_LOGD("End to define %s pass pattern.", pass_name5.c_str());

  string pass_name6 = "TbeDepthwiseConvDequantPowerFusion";
  BufferFusionPattern *pattern6 = new (std::nothrow) BufferFusionPattern(pass_name6);
  FE_CHECK((pattern6 == nullptr),
           REPORT_FE_ERROR("[SubGraphOpt][DepthwiseConvDeQntFus][DefPtn] new an object failed."), return patterns);
  FE_LOGD("Start to define %s pass pattern.", pass_name6.c_str());
  // define pattern rules Convolution-->AcendDeQuant
  pattern6->AddOpDesc(PATTERN_DEQUANT, {OP_PATTERN_DEQUANT}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_CONV, {OP_PATTERN_DEPTHWISE_CONV}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_POWER1, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_RELU6, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_POWER2, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_ELEMWISE, {OP_PATTERN_ELEMWISE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_QUANT, {OP_PATTERN_QUANT}, TBE_PATTERN_NUM_NONE, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .AddOpDesc(PATTERN_OTHER_INPUT2, {TBE_PATTERN_INPUT_NODE}, TBE_PATTERN_NUM_DEFAULT, TBE_PATTERN_NUM_DEFAULT)
      .SetHead({PATTERN_CONV})
      .SetOutputs(PATTERN_CONV, {PATTERN_DEQUANT})
      .SetOutputs(PATTERN_DEQUANT, {PATTERN_POWER1, PATTERN_ELEMWISE}, TBE_OUTPUT_BRANCH_MULTI)
      .SetOutputs(PATTERN_POWER1, {PATTERN_RELU6})
      .SetOutputs(PATTERN_RELU6, {PATTERN_POWER2})
      .SetOutputs(PATTERN_POWER2, {PATTERN_ELEMWISE})
      .SetOutputs(PATTERN_ELEMWISE, {PATTERN_QUANT}, TBE_OUTPUT_BRANCH_SINGE, IGNORE_INPUT_NUM, IGNORE_OUTPUT_NUM)
      .SetOutputs(PATTERN_OTHER_INPUT, {PATTERN_DEQUANT});
  patterns.push_back(pattern6);
  FE_LOGD("End to define %s pass pattern.", pass_name6.c_str());


  return patterns;
}

void SearchMatchNode(const BufferFusionMapping &mapping,
                     vector<ge::NodePtr> &fusion_nodes) {
    for (auto &item : mapping) {
        auto opdesc = find(item.first->types.begin(), item.first->types.end(), TBE_PATTERN_OUTPUT_NODE);

    if (opdesc != item.first->types.end()) {
        for (auto &node : item.second) {
            auto node_ptr = find(fusion_nodes.begin(), fusion_nodes.end(), node);
            fusion_nodes.erase(node_ptr);
            }
        }
    }
}

/*
 * @brief: parse nodes matched in mapping and call DoFusion
 * @param [in] graph: original graph
 * @param [out] mapping: nodes matched by pattern
 * @return bool: fusion status ok or not.
 */
Status TbeDepthwiseConvDequantFusionPass::GetFusionNodes(const BufferFusionMapping &mapping,
                                                         vector<ge::NodePtr> &fusion_nodes) {
  FE_LOGD("Begin to do TbeDepthwiseConvDequantFusion!");
  vector<ge::NodePtr> elem2_node = GetMatchedNodesByDescName(PATTERN_MUL, mapping);
  vector<ge::NodePtr> elem3_node = GetMatchedNodesByDescName(PATTERN_SIGMOID, mapping);
  vector<ge::NodePtr> relu_node = GetMatchedNodesByDescName(PATTERN_RELU, mapping);
  vector<ge::NodePtr> power_node = GetMatchedNodesByDescName(PATTERN_POWER1, mapping);
  vector<ge::NodePtr> eletwise_node = GetMatchedNodesByDescName(PATTERN_ELEMWISE, mapping);
  if (!relu_node.empty()) {
    if ((relu_node[0]->GetType() != "LeakyRelu") && (relu_node[0]->GetType() != "Relu6") &&
        (relu_node[0]->GetType() != "Relu6D")) {
      FE_LOGI(
          "The optype of node[%s] should be LeakyRelu or Relu6 or Relu6D,"
          "but actually is [%s], no need to do fusion.",
          relu_node[0]->GetName().c_str(), relu_node[0]->GetType().c_str());
      return SUCCESS;
    }
  }
  if (!elem2_node.empty()) {
    if (elem2_node[0]->GetType() != "Mul") {
      FE_LOGI(
          "The optype of node[%s] should be Mul,"
          "but actually is [%s], no need to do fusion.",
          elem2_node[0]->GetName().c_str(), elem2_node[0]->GetType().c_str());
      return SUCCESS;
    }
  }
  if (!elem3_node.empty()) {
    if (elem3_node[0]->GetType() != "Sigmoid") {
      FE_LOGI(
          "The optype of node[%s] should be Sigmoid,"
          "but actually is [%s], no need to do fusion.",
          elem3_node[0]->GetName().c_str(), elem3_node[0]->GetType().c_str());
      return SUCCESS;
    }
    FE_LOGD("DepthwiseConvElemwise + dequant + sigmoid + mul + quant ub fusion!");
  }
  if (!power_node.empty()) {
    if (power_node[0]->GetType() != "Power") {
      FE_LOGI(
          "The optype of node[%s] should be Sigmoid,"
          "but actually is [%s], no need to do fusion.",
          power_node[0]->GetName().c_str(), power_node[0]->GetType().c_str());
      return SUCCESS;
    }
  }
  if (!eletwise_node.empty()) {
    if (eletwise_node[0]->GetType() != "Eltwise") {
      FE_LOGI(
          "The optype of node[%s] should be Sigmoid,"
          "but actually is [%s], no need to do fusion.",
          eletwise_node[0]->GetName().c_str(), eletwise_node[0]->GetType().c_str());
      return SUCCESS;
    }
  }
  fusion_nodes = GetMatchedNodes(mapping);
  // the output_data can't be fused
  SearchMatchNode(mapping, fusion_nodes);

  FE_LOGD("End to do TbeDepthwiseConvDequantFusion!");
  return SUCCESS;
}
}  // namespace fe
